"""
安全中间件 - 提供安全头部和HTTPS重定向
"""
from fastapi import Request, Response, HTTPException, status
from fastapi.responses import RedirectResponse, JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
from urllib.parse import urlparse
import os
import secrets
import time
from typing import Optional
from .config import get_settings


class SecurityMiddleware(BaseHTTPMiddleware):
    """安全中间件"""
    
    def __init__(self, app):
        super().__init__(app)
        self.settings = get_settings()
        self.csrf_tokens = {}  # 存储CSRF令牌
    
    async def dispatch(self, request: Request, call_next):
        # HTTPS强制重定向
        if self.settings.force_https and request.url.scheme != "https":
            # 在生产环境中强制HTTPS
            if not self._is_local_development(request):
                url = request.url.replace(scheme="https")
                return RedirectResponse(url=str(url))
        
        # CSRF保护
        if self._should_check_csrf(request):
            if not await self._validate_csrf_token(request):
                return JSONResponse(
                    status_code=status.HTTP_403_FORBIDDEN,
                    content={"detail": "CSRF token validation failed"}
                )
        
        response = await call_next(request)
        
        # 添加安全头部
        self._add_security_headers(response)
        
        # 为表单页面添加CSRF令牌
        if self._should_add_csrf_token(request, response):
            self._add_csrf_token_to_response(response)
        
        return response
    
    def _is_local_development(self, request: Request) -> bool:
        """检查是否为本地开发环境"""
        host = request.headers.get("host", "")
        return (
            host.startswith("localhost") or
            host.startswith("127.0.0.1") or
            host.startswith("192.168.") or
            host.startswith("10.") or
            host.endswith(".local")
        )
    
    def _add_security_headers(self, response: Response):
        """添加安全头部"""
        # HSTS
        if self.settings.security_hsts_enabled:
            hsts_value = f"max-age={self.settings.security_hsts_max_age}"
            if self.settings.security_hsts_include_subdomains:
                hsts_value += "; includeSubDomains"
            if self.settings.security_hsts_preload:
                hsts_value += "; preload"
            response.headers["Strict-Transport-Security"] = hsts_value
        
        # CSP
        if self.settings.security_csp_enabled:
            csp_parts = []
            
            if self.settings.security_csp_default_src:
                csp_parts.append(f"default-src {' '.join(self.settings.security_csp_default_src)}")
            
            if self.settings.security_csp_script_src:
                csp_parts.append(f"script-src {' '.join(self.settings.security_csp_script_src)}")
            
            if self.settings.security_csp_style_src:
                csp_parts.append(f"style-src {' '.join(self.settings.security_csp_style_src)}")
            
            if self.settings.security_csp_img_src:
                csp_parts.append(f"img-src {' '.join(self.settings.security_csp_img_src)}")
            
            if self.settings.security_csp_font_src:
                csp_parts.append(f"font-src {' '.join(self.settings.security_csp_font_src)}")
            
            if self.settings.security_csp_connect_src:
                csp_parts.append(f"connect-src {' '.join(self.settings.security_csp_connect_src)}")
            
            if self.settings.security_csp_frame_src:
                csp_parts.append(f"frame-src {' '.join(self.settings.security_csp_frame_src)}")
            
            if self.settings.security_csp_object_src:
                csp_parts.append(f"object-src {' '.join(self.settings.security_csp_object_src)}")
            
            response.headers["Content-Security-Policy"] = "; ".join(csp_parts)
        
        # 其他安全头部
        if self.settings.security_x_frame_options:
            response.headers["X-Frame-Options"] = self.settings.security_x_frame_options
        
        if self.settings.security_x_content_type_options:
            response.headers["X-Content-Type-Options"] = self.settings.security_x_content_type_options
        
        if self.settings.security_x_xss_protection:
            response.headers["X-XSS-Protection"] = self.settings.security_x_xss_protection
        
        if self.settings.security_referrer_policy:
            response.headers["Referrer-Policy"] = self.settings.security_referrer_policy
        
        if self.settings.security_permissions_policy:
            response.headers["Permissions-Policy"] = self.settings.security_permissions_policy
        
        # 移除服务器信息
        response.headers["Server"] = "MinerU-Doc-Service"
        
        # 防止MIME类型嗅探
        response.headers["X-Download-Options"] = "noopen"
        
        # 防止点击劫持
        response.headers["X-Permitted-Cross-Domain-Policies"] = "none"
    
    def _should_check_csrf(self, request: Request) -> bool:
        """检查是否需要进行CSRF验证"""
        # 只对非GET、HEAD、OPTIONS方法的请求进行CSRF验证
        if request.method in ["GET", "HEAD", "OPTIONS"]:
            return False
        
        # 检查是否为API请求
        if request.url.path.startswith("/api/"):
            # 对于登录API，不进行CSRF验证（允许未认证访问）
            if request.url.path in ["/api/v1/auth/login", "/api/v1/auth/register"]:
                return False
            
            # 对于其他API请求，检查是否有Authorization头
            auth_header = request.headers.get("Authorization")
            if auth_header and auth_header.startswith("Bearer "):
                return False  # 已认证的API请求不需要CSRF验证
        
        return True
    
    async def _validate_csrf_token(self, request: Request) -> bool:
        """验证CSRF令牌"""
        # 从Cookie中获取CSRF令牌
        csrf_token_cookie = request.cookies.get("csrf_token")
        if not csrf_token_cookie:
            return False
        
        # 从请求头获取CSRF令牌
        csrf_token_request = (
            request.headers.get("X-CSRF-Token") or
            request.headers.get("X-XSRF-Token")
        )
        
        # 如果没有在请求头中找到，尝试从表单数据中获取
        if not csrf_token_request and request.headers.get("content-type", "").startswith("multipart/form-data"):
            form_data = await request.form()
            csrf_token_request = form_data.get("csrf_token")
        
        if not csrf_token_request:
            return False
        
        # 验证令牌是否匹配
        return secrets.compare_digest(csrf_token_cookie, csrf_token_request)
    
    def _should_add_csrf_token(self, request: Request, response: Response) -> bool:
        """检查是否应该添加CSRF令牌"""
        # 只对HTML页面添加CSRF令牌
        content_type = response.headers.get("content-type", "")
        return "text/html" in content_type
    
    def _add_csrf_token_to_response(self, response: Response):
        """为响应添加CSRF令牌"""
        if hasattr(response, "body") and response.body:
            try:
                body_str = response.body.decode()
                if "</form>" in body_str:
                    # 生成CSRF令牌
                    csrf_token = secrets.token_urlsafe(32)
                    
                    # 将CSRF令牌添加到响应中
                    csrf_input = f'<input type="hidden" name="csrf_token" value="{csrf_token}">'
                    body_str = body_str.replace("</form>", f"{csrf_input}</form>")
                    
                    # 更新响应体
                    response.body = body_str.encode()
                    
                    # 设置CSRF令牌Cookie
                    response.set_cookie(
                        key="csrf_token",
                        value=csrf_token,
                        httponly=False,
                        secure=True,
                        samesite="strict",
                        max_age=3600  # 1小时
                    )
            except Exception:
                # 如果解析失败，忽略错误
                pass